#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   __init__.py
@Time    :   2024-07-09 12:57:20
@Author  :   chakcy 
@Email   :   947105045@qq.com
@description   :   jwt中间件
'''
from typing import Any

from starlette.middleware.base import BaseHTTPMiddleware, DispatchFunction
from starlette.types import ASGIApp
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse

from app.dependencies.jwt import decode_token
from .is_pass import isPass

class JWTMiddleware(BaseHTTPMiddleware):
    """JWT 中间件"""
    def __init__(self, app:ASGIApp, header_namespace:str):
        super().__init__(app)
        self.header_namespace = header_namespace

    async def dispatch(self, # type: ignore
                       request: Request, 
                       call_next: DispatchFunction) -> Any: 
        try:
            if isPass(request.url.path):
                response: Any = await call_next(request) # type: ignore
                return response
            
            token = request.headers.get(self.header_namespace)
            if token is None:
                raise HTTPException(  
                    status_code=status.HTTP_401_UNAUTHORIZED,  
                    detail="Authorization header is missing"  
                )
            
            if not token.startswith("Bearer "):
                raise HTTPException(  
                    status_code=status.HTTP_401_UNAUTHORIZED,  
                    detail="Authorization header must start with Bearer"  
                )

            data = decode_token(token)

            if data is None:
                raise HTTPException(  
                    status_code=status.HTTP_401_UNAUTHORIZED,  
                    detail="Invalid token"  
                )
            
            request.state.data = data
            request.state.token = token
            response: Any = await call_next(request) # type: ignore
            return response
        
        except HTTPException as e:
            return JSONResponse(status_code=e.status_code, content={"detail":e.detail})
        